import numpy as np
import pandas as pd
from config import INITIAL_CASH, TRANSACTION_COST, MAX_POSITION_RATIO, WINDOW_SIZE, ASSET_CODES


class TradingEnv:
    def __init__(self, data):
        self.data = data  # 预处理后的多资产数据
        self.asset_types = list(ASSET_CODES.keys())  # 资产类型列表
        self.n_assets = len(self.asset_types)  # 资产数量
        self.n_features = len(data.columns)  # 特征数量
        self.max_steps = len(data) - WINDOW_SIZE - 1  # 最大步数

        # 环境状态初始化
        self.reset()

    def reset(self):
        """重置环境"""
        self.current_step = WINDOW_SIZE  # 当前步骤（从窗口结束后开始）
        self.cash = INITIAL_CASH  # 现金余额
        self.positions = np.zeros(self.n_assets)  # 各资产持仓数量
        self.total_asset = INITIAL_CASH  # 总资产（现金+持仓市值）
        self.history = []  # 历史记录（用于回测）
        return self._get_state()

    def _get_state(self):
        """获取当前状态：过去WINDOW_SIZE天的特征数据"""
        start = self.current_step - WINDOW_SIZE
        end = self.current_step
        state = self.data.iloc[start:end].values  # (WINDOW_SIZE, n_features)
        return state.flatten()  # 展平为一维数组

    def _calculate_reward(self, prev_total_asset):
        """计算奖励：收益-风险惩罚"""
        # 收益奖励（当前总资产变化率）
        profit_reward = (self.total_asset - prev_total_asset) / prev_total_asset

        # 风险惩罚（波动率惩罚）
        if len(self.history) > 10:
            recent_returns = [h["total_asset"] for h in self.history[-10:]]
            volatility = np.std(recent_returns) / np.mean(recent_returns)
            risk_penalty = -volatility * 0.4  # 惩罚系数0.1
        else:
            risk_penalty = 0

        return profit_reward + risk_penalty

    def step(self, action):
        """执行动作：0=持有，1=买入，2=卖出（对单资产）"""
        # 动作解析：action为数组，每个元素对应一个资产的动作（0/1/2）
        prev_total_asset = self.total_asset
        current_prices = self._get_current_prices()  # 当前各资产收盘价

        # 执行每个资产的动作
        for i in range(self.n_assets):
            asset_type = self.asset_types[i]
            price = current_prices[i]
            action_i = action[i]

            if action_i == 1:  # 买入
                # 计算可买入最大数量（不超过最大仓位）
                max_buy_amount = (self.cash * MAX_POSITION_RATIO) // price
                if max_buy_amount > 0:
                    cost = max_buy_amount * price * (1 + TRANSACTION_COST)
                    self.cash -= cost
                    self.positions[i] += max_buy_amount

            elif action_i == 2:  # 卖出
                # 卖出全部持仓
                if self.positions[i] > 0:
                    revenue = self.positions[i] * price * (1 - TRANSACTION_COST)
                    self.cash += revenue
                    self.positions[i] = 0

        # 更新总资产
        position_value = np.sum(self.positions * current_prices)
        self.total_asset = self.cash + position_value

        # 记录历史
        self.history.append({
            "step": self.current_step,
            "cash": self.cash,
            "positions": self.positions.copy(),
            "position_value": position_value,
            "total_asset": self.total_asset
        })

        # 计算奖励和终止状态
        reward = self._calculate_reward(prev_total_asset)
        self.current_step += 1
        done = self.current_step >= self.max_steps

        # 信息输出
        info = {
            "total_asset": self.total_asset,
            "current_step": self.current_step
        }

        return self._get_state(), reward, done, info

    def _get_current_prices(self):
        """获取当前各资产收盘价"""
        prices = []
        for asset_type in self.asset_types:
            price_col = f"close_{asset_type}"
            prices.append(self.data.iloc[self.current_step][price_col])
        return np.array(prices)

    def render(self):
        """打印当前状态"""
        print(f"Step: {self.current_step}, Total Asset: {self.total_asset:.2f}, Cash: {self.cash:.2f}")